import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

'''
使用PyTorch库来创建和训练全连接神经网络以拟合抛物线的Python代码示例。
生成了一些带有噪声的训练数据。
定义了一个包含两个隐藏层（每层64个神经元）的全连接神经网络类。
实例化了模型，并定义了损失函数（均方误差）和优化器（Adam）。
训练了模型100个周期。
使用训练好的模型对测试数据进行预测。
可视化了训练数据、神经网络的预测和真实的抛物线。
请注意，由于我们添加了噪声，神经网络的预测可能不会完美地匹配真实的抛物线，但它应该能够捕捉到大致的趋势。此外，你可能需要调整网络的结构（如层的数量、每层的神经元数量）或训练参数（如学习率、训练周期数）以获得更好的性能。
'''

# 定义全连接神经网络模型
class FullyConnectedNN(nn.Module):
    def __init__(self):
        super(FullyConnectedNN, self).__init__()
        self.fc1 = nn.Linear(1, 64)  # 输入层和第一个隐藏层
        self.fc2 = nn.Linear(64, 64)  # 第二个隐藏层
        self.fc3 = nn.Linear(64, 1)  # 输出层
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x) # 输出层没有激活函数
        return x


# 生成训练数据
np.random.seed(0)  # 为了结果可重复
x_train_1 = np.linspace(0.1, 1, 5).reshape(-1, 1)
x_train_2 = np.linspace(2, 3, 4).reshape(-1, 1)
x_train_7 = np.linspace(7, 8, 4).reshape(-1, 1)
x_train = np.concatenate((x_train_1, x_train_2, x_train_7))
y_train = x_train ** 2 + np.random.normal(0, 10, x_train.shape)  # 添加一些噪声

# 转换为PyTorch张量
x_train_tensor = torch.tensor(x_train, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train, dtype=torch.float32).reshape(-1, 1)  # 转换为列向量

# 实例化模型、定义损失函数和优化器
model = FullyConnectedNN()
criterion = nn.MSELoss()  # 使用均方误差作为损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)  # 使用Adam优化器

# 训练模型
num_epochs = 1500
for epoch in range(num_epochs):
    model.train()  # 设置模型为训练模式
    optimizer.zero_grad()  # 清空梯度
    outputs = model(x_train_tensor)  # 前向传播
    loss = criterion(outputs, y_train_tensor)  # 计算损失
    loss.backward()  # 反向传播
    optimizer.step()  # 更新参数

    if (epoch + 1) % 100 == 0:
        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 测试模型（使用训练数据进行测试，实际应用中应使用未见过的数据）
model.eval()  # 设置模型为评估模式
x_test = np.linspace(-10, 10, 100).reshape(-1, 1)  # 测试数据
x_test_tensor = torch.tensor(x_test, dtype=torch.float32)
y_pred = model(x_test_tensor).detach().numpy()  # 预测并转换为NumPy数组

# 可视化结果
plt.scatter(x_train, y_train, label='Training Data')
plt.plot(x_test, y_pred, color='red', label='NN Prediction')
plt.plot(x_test, x_test ** 2, color='green', label='True Parabola')
plt.legend()
plt.show()